from nlp_tools.tasks.abs_task_model import ABCTaskModel
from tensorflow.keras import backend as K
from collections import Counter
import os

import pandas as pd

df_test = pd.read_csv("/home/qiufengfeng/nlp/competition/datagrand/基于大规模预训练模型的风险事件标签识别/datagrand_2021_test.csv")
train_model_path = './tf_models'


def load_model_list(model_path):
    model_list = []
    for file in os.listdir(model_path):
        model_list.append(os.path.join(model_path,file))
    return model_list

def voteing(predict_list):
    result = []
    for i in range(len(predict_list[0])):
        count = []
        predict_label = [item[i] for item in predict_list]
        predict_label_dict = dict(Counter(predict_label))
        sorted_dict = sorted(predict_label_dict.items(), key=lambda x: x[1], reverse=True)
        result.append(sorted_dict[0][0])
    return result




predict_list = []
model_list = load_model_list(train_model_path)

for model_path in model_list:
    K.clear_session()
    model = ABCTaskModel.load_model(model_path)
    y_pred = model.predict(df_test['text'].tolist())
    predict_list.append(y_pred)
finnal_result = voteing(predict_list)


pd.DataFrame({"id": df_test['id'], "label": finnal_result}).to_csv("submission.csv", index=False)